import os
import logging
from typing import Optional, Dict, Any, List
from contextlib import AsyncExitStack

from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
from dotenv import load_dotenv

# 配置日志
logger = logging.getLogger(__name__)


class MCPClient:
    """
    MCP客户端类，用于连接并调用MCP服务器中的工具
    """
    def __init__(self):
        self.session: Optional[ClientSession] = None
        self.exit_stack = AsyncExitStack()
        self.connected = False
        logger.debug("MCPClient实例已创建")
    
    async def connect_to_server(self, command: str, args: List[str], env: Dict[str, str] = None):
        """
        连接到MCP服务器
        
        参数:
            command: 启动服务器的命令
            args: 命令行参数
            env: 环境变量
        """
        if self.connected:
            logger.warning("已经连接到服务器，请先断开连接")
            return

        logger.info(f"正在连接到MCP服务器，命令: {command}，参数: {args}")
        
        # 为stdio连接创建服务器参数
        server_params = StdioServerParameters(
            command=command,
            args=args,
            env=env
        )

        try:
            # 创建stdio客户端
            logger.debug("创建stdio客户端")
            stdio_transport = await self.exit_stack.enter_async_context(
                stdio_client(server_params)
            )
            stdio, write = stdio_transport
            
            # 创建ClientSession对象
            logger.debug("创建ClientSession对象")
            self.session = await self.exit_stack.enter_async_context(
                ClientSession(stdio, write)
            )
            
            # 初始化ClientSession
            logger.debug("初始化ClientSession")
            await self.session.initialize()
            
            self.connected = True
            logger.info("成功连接到MCP服务器")
        except Exception as e:
            logger.error(f"连接服务器失败: {str(e)}", exc_info=True)
            await self.cleanup()
            
    async def list_tools(self):
        """列出服务器上可用的工具"""
        if not self.is_connected():
            logger.warning("未连接到服务器，请先连接")
            return None
        
        try:
            logger.debug("正在获取工具列表")
            response = await self.session.list_tools()
            logger.info(f"获取到 {len(response.tools)} 个工具")
            return response.tools
        except Exception as e:
            logger.error(f"获取工具列表失败: {str(e)}", exc_info=True)
            return None
    
    async def call_tool(self, tool_name: str, params: Dict[str, Any]):
        """
        调用MCP服务器上的工具
        
        参数:
            tool_name: 工具名称
            params: 工具参数
        """
        if not self.is_connected():
            logger.warning("未连接到服务器，请先连接")
            return None
        
        try:
            logger.info(f"正在调用工具 {tool_name}，参数: {params}")
            result = await self.session.call_tool(tool_name, params)
            logger.debug(f"工具 {tool_name} 调用成功，返回结果: {result}")
            return result
        except Exception as e:
            logger.error(f"调用工具 {tool_name} 失败: {str(e)}", exc_info=True)
            return None
    
    def is_connected(self):
        """检查是否已连接到服务器"""
        return self.connected and self.session is not None
    
    async def cleanup(self):
        """清理资源"""
        logger.info("正在清理资源")
        self.connected = False
        await self.exit_stack.aclose()
        self.session = None
        logger.debug("资源清理完成")


def convert_mcp_tools_to_openai_format(mcp_tools):
    """
    将MCP工具转换为OpenAI工具格式
    
    参数:
        mcp_tools: MCP工具列表
        
    返回:
        符合OpenAI格式的工具列表
    """
    if not mcp_tools:
        return []
    
    logger.debug("开始转换MCP工具为OpenAI格式")
    openai_tools = []
    
    for tool in mcp_tools:
        openai_tool = {
            "type": "function",
            "function": {
                "name": tool.name,
                "description": tool.description,
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        }
        
        # 处理参数
        if hasattr(tool, 'parameters') and tool.parameters:
            for param_name, param_info in tool.parameters.items():
                openai_tool["function"]["parameters"]["properties"][param_name] = {
                    "type": param_info.get("type", "string"),
                    "description": param_info.get("description", "")
                }
                
                # 如果参数是必需的
                if param_info.get("required", False):
                    openai_tool["function"]["parameters"]["required"].append(param_name)
        
        openai_tools.append(openai_tool)
    
    logger.debug(f"转换完成，共 {len(openai_tools)} 个工具")
    return openai_tools


# 在模块级别配置基本的日志处理，如果没有外部配置
def setup_default_logging():
    """设置默认的日志配置，如果没有外部配置时使用"""
    if not logging.root.handlers:  # 只有在没有处理器时才配置
        logging.basicConfig(
            level=logging.INFO,
            format='%(asctime)s [%(levelname)s] %(name)s:%(lineno)d - %(message)s',
            datefmt='%Y-%m-%d %H:%M:%S'
        )
        logger.debug("已设置默认日志配置")

# 调用默认日志配置
setup_default_logging()

